from s4dd import S4forDenovoDesign
import torch

# 主程序入口
if __name__ == "__main__":
    # 初始化 S4 模型
    s4 = S4forDenovoDesign(
        n_max_epochs=1,  # 每次只训练 1 个 epoch，手动控制循环
        batch_size=128,  # 批量大小
        device="cuda" if torch.cuda.is_available() else "cpu",  # 自动检测是否使用 GPU
    )

    # 数据集路径
    training_molecules_path = "./datasets/chemblv31/train_70.txt"
    val_molecules_path = "./datasets/chemblv31/valid.zip"
    finetune_training_molecules_path = "./datasets/pkm2/train.zip"
    finetune_val_molecules_path = "./datasets/pkm2/valid.zip"

    total_epochs = 5  # 预训练和微调的总轮数

    print("开始模型预训练和微调...")

    for epoch in range(1, total_epochs + 1):
        # --- 预训练 ---
        print(f"开始第 {epoch} 轮预训练...")
        s4.train(training_molecules_path=training_molecules_path, val_molecules_path=val_molecules_path)

        # 保存当前预训练模型
        pretrain_model_path = f"./new_models/s4_pretrained_epoch_{epoch}"
        s4.save(pretrain_model_path)
        print(f"第 {epoch} 轮预训练模型已保存到 {pretrain_model_path}！")

        # --- 微调 ---
        print(f"开始第 {epoch} 轮微调...")
        s4.train(training_molecules_path=finetune_training_molecules_path,
                 val_molecules_path=finetune_val_molecules_path)

        # 保存当前微调模型
        finetune_model_path = f"./new_models/s4_finetuned_epoch_{epoch}"
        s4.save(finetune_model_path)
        print(f"第 {epoch} 轮微调模型已保存到 {finetune_model_path}！")

    print("预训练和微调过程完成！")

    # 保存最终预训练和微调模型
    final_pretrained_model_path = "./new_models/s4_pretrained_full"
    final_finetuned_model_path = "./new_models/s4_finetuned_final"
    s4.save(final_pretrained_model_path)
    print(f"最终预训练模型已保存到 {final_pretrained_model_path}！")
    s4.save(final_finetuned_model_path)
    print(f"最终微调模型已保存到 {final_finetuned_model_path}！")


    # print("开始设计新的分子...")
    # designs, lls = s4.design_molecules(n_designs=100, batch_size=32, temperature=1.0)  # 设计 100 个新分子
    # print("分子设计完成！")
    #
    #
    # output_file_path = "./generated_molecules_full.txt"
    # with open(output_file_path, "w") as f:
    #     for i, molecule in enumerate(designs):
    #         f.write(f"Molecule {i + 1}: {molecule}\n")
    # print(f"生成的分子已保存到 {output_file_path}！")
    #
    #
    # print("生成的分子设计如下：")
    # for i, molecule in enumerate(designs):
    #     print(f"Molecule {i + 1}: {molecule}")
